# Copyright (c) 2022 Presto Labs Pte. Ltd.
# Author: lucasyoo

import os
import json
from cc.appcoin2.strategy.hamm.support.util import StratConfigHamm
from cc.appcoin2.strategy.vmm.support.util import StratConfigVmm

from absl import app, flags

FLAGS = flags.FLAGS


def gen_driver_and_config():
  os.makedirs(f'{FLAGS.sim_result_dir}/configs', exist_ok=True)
  os.makedirs(f'{FLAGS.sim_result_dir}/drivers', exist_ok=True)

  base_input = {
    'config_path': FLAGS.config_path,
    'is_sim': FLAGS.is_sim,
    'interval': [FLAGS.start_dt, FLAGS.end_dt],
    'sim_dir': FLAGS.tmp_dirname,
    'analyzer_feed_machine': FLAGS.analyzer_feed_machine,
  }

  if FLAGS.strategy_group == 'vmm':
    base_input.update({
        'is_approach': FLAGS.is_approach,
        'initial_pos': FLAGS.initial_pos,
        'include_tag_info': FLAGS.include_tag_info
    })
    strat_config = StratConfigVmm
  elif FLAGS.strategy_group == 'hamm':
    strat_config = StratConfigHamm
  else:
    raise ValueError(f'unknown strategy_group: {FLAGS.strategy_group}')

  strat_config_cls = strat_config(**base_input)
  strat_config_cls.update_params(json.loads(FLAGS.params))
  strat_config_cls.gen_config()
  with open(FLAGS.driver_file_name, 'w') as fo:
    json.dump(strat_config_cls.app_config, fo, indent=2)

  with open(FLAGS.config_file_name, 'w') as fo:
    json.dump(strat_config_cls.config, fo, indent=2)


def main(_):
  gen_driver_and_config()


if __name__ == '__main__':
  flags.DEFINE_string('analyzer_feed_machine', None, '')
  flags.DEFINE_string('config_path', '', 'config_path')
  flags.DEFINE_bool('is_sim', True, 'is_sim')
  flags.DEFINE_string('driver_file_name', '', 'driver file name')
  flags.DEFINE_string('config_file_name', '', 'config file name')
  flags.DEFINE_string('start_dt', '', 'start_dt')
  flags.DEFINE_string('end_dt', '', 'end_dt')
  flags.DEFINE_string('params', '', 'params')
  flags.DEFINE_string('tmp_dirname', '', 'tmp_dirname')
  flags.DEFINE_string('sim_result_dir', '', 'sim_result_dir')
  flags.DEFINE_string('strategy_group', '', 'strategy_group')
  flags.DEFINE_bool('is_approach', False, 'is_approach')
  flags.DEFINE_float('initial_pos', 0.0, 'initial_pos')
  flags.DEFINE_bool('include_tag_info', False, 'include_tag_info')

  app.run(main)
